#include "evdeploy/cv/common/common.h"


namespace ev
{
namespace vision
{

float IOU(const cv::Rect &b1, const cv::Rect &b2)
{
    auto intersec = b1 & b2;
    return static_cast<float>(intersec.area()) / (b1.area() + b2.area() - intersec.area());
}

void NMS(std::vector<BoxInfo> &objects, float iou_thresh, bool class_aware)
{
    auto cmp_lammda = [](const BoxInfo &b1, const BoxInfo &b2)
    { return b1.score < b2.score; };
    std::sort(objects.begin(), objects.end(), cmp_lammda);
    for (int i = 0; i < objects.size(); ++i)
    {
        for (int j = i + 1; j < objects.size(); ++j)
        {
            cv::Rect rect1 = cv::Rect{objects[i].x1, objects[i].y1, objects[i].x2 - objects[i].x1, objects[i].y2 - objects[i].y1};
            cv::Rect rect2 = cv::Rect{objects[j].x1, objects[j].y1, objects[j].x2 - objects[j].x1, objects[j].y2 - objects[j].y1};
            if (IOU(rect1, rect2) > iou_thresh && (class_aware == false || (class_aware == true && objects[i].label == objects[j].label)))
            {
                objects[i].score = 0.f;
            }
        }
    }
    auto iter = objects.begin();
    while (iter != objects.end())
    {
        if (iter->score < 0.000001)
        {
            iter = objects.erase(iter);
        }
        else
        {
            ++iter;
        }
    }
}

}//namespace vision
}//namepspace ev